package cn.lxb.emall.gw.filter;

import java.util.List;
import java.util.Set;
import java.util.function.Function;

import org.reactivestreams.Publisher;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.rewrite.MessageBodyDecoder;
import org.springframework.cloud.gateway.filter.factory.rewrite.MessageBodyEncoder;
import org.springframework.cloud.gateway.filter.factory.rewrite.ModifyResponseBodyGatewayFilterFactory;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.codec.HttpMessageReader;
import org.springframework.web.server.ServerWebExchange;

import lombok.Getter;
import lombok.Setter;
import reactor.core.publisher.Mono;

/** 可根据条件修改响应体的过滤器工厂 */
public class ConditionalModifyResponseBodyGatewayFilterFactory extends ModifyResponseBodyGatewayFilterFactory {

	@Setter
	@Getter
	Function<ServerWebExchange, Boolean> condition;

	public ConditionalModifyResponseBodyGatewayFilterFactory(List<HttpMessageReader<?>> messageReaders,
			Set<MessageBodyDecoder> messageBodyDecoders, Set<MessageBodyEncoder> messageBodyEncoders) {
		super(messageReaders, messageBodyDecoders, messageBodyEncoders);
	}

	@Override
	public GatewayFilter apply(Config config) {
		return (exchange, chain) -> chain.filter(exchange.mutate()
				.response(new ConditionalModifiedServerHttpResponse(exchange, config, condition)).build());
	}

	public class ConditionalModifiedServerHttpResponse extends ModifiedServerHttpResponse {
		private final ServerWebExchange exchange;
		Function<ServerWebExchange, Boolean> condition;

		public ConditionalModifiedServerHttpResponse(ServerWebExchange exchange, Config config,
				Function<ServerWebExchange, Boolean> condition) {
			super(exchange, config);
			this.exchange = exchange;
			this.condition = condition;
		}

		@Override
		public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
			if (Boolean.TRUE.equals(condition.apply(exchange))) {
				return super.writeWith(body);
			}
			return getDelegate().writeWith(body);
		}
	}

}
